import torch.nn as nn
import torch.nn.functional as F

import netRes


class Net(nn.Module):
    def __init__(self, n_blocks=10,res_chans = 32):
        super().__init__()
        self.conv1 = nn.Conv2d(1,res_chans,kernel_size=3,stride=1,padding=1)
        self.resblock = nn.Sequential(
            *(n_blocks *
              [netRes.NetRes(res_chans = res_chans)]))
        self.flatten = nn.Flatten()
        self.liner1 = nn.Linear(14 * 14 * res_chans, 256)
        self.liner2 = nn.Linear(256, 10)

    def forward(self, x):
        out = F.max_pool2d(F.relu(self.conv1(x)),2)
        out = self.resblock(out)

        out = self.flatten(out)
        out = F.relu(self.liner1(out))
        out = F.relu(self.liner2(out))
        out = F.log_softmax(out, dim=1)
        return out